"""
股票K线数据获取模块
支持从多个数据源获取股票和期货的K线数据
包括 Yahoo Finance (国际股票), AKShare (A股/期货) 和 Tushare (专业金融数据)
"""

import pandas as pd
import logging
import time
import os
from datetime import datetime, timedelta
from config import (
    DATA_SOURCE, 
    PERIOD_MAPPING, 
    API_TIMEOUT, 
    TUSHARE_TOKEN,
    SUPPORTED_PERIODS,
    ENABLE_CACHE,
    CACHE_EXPIRY_HOURS,
    MAX_RETRIES,
    RETRY_DELAY,
    OUTPUT_FILENAME_FORMAT,
    OUTPUT_ENCODING,
    LOG_LEVEL
)

# 设置日志
logger = logging.getLogger(__name__)
logger.setLevel(LOG_LEVEL)

class DataFetcher:
    def __init__(self):
        """
        初始化数据获取器
        """
        self.data_source = DATA_SOURCE
        self.cache_dir = "data_cache"
        self._create_cache_dir()
        
        # 初始化数据源API
        if self.data_source == 'tushare' and TUSHARE_TOKEN:
            import tushare as ts
            ts.set_token(TUSHARE_TOKEN)
            self.tushare_pro = ts.pro_api()
            logger.info("Tushare API初始化完成")
        elif self.data_source == 'baostock':
            try:
                import baostock as bs
                self.bs = bs
                logger.info("Baostock数据源初始化完成")
            except ImportError:
                logger.error("Baostock模块未安装，请运行: pip install baostock")
                raise
        else:
            if self.data_source == 'tushare':
                logger.warning("Tushare token未配置，将无法使用Tushare数据源")

    def _create_cache_dir(self):
        """创建缓存目录"""
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
            logger.info(f"创建缓存目录: {self.cache_dir}")

    def _get_cache_filename(self, symbol, period_name):
        """
        生成缓存文件名
        :param symbol: 股票代码
        :param period_name: K线周期名称
        :return: 缓存文件路径
        """
        safe_symbol = symbol.replace(':', '_').replace('/', '_')
        return os.path.join(self.cache_dir, f"{safe_symbol}_{period_name}.csv")

    def _is_cache_valid(self, cache_file):
        """
        检查缓存是否有效
        :param cache_file: 缓存文件路径
        :return: 布尔值，表示缓存是否有效
        """
        if not os.path.exists(cache_file):
            return False
            
        # 检查缓存时间
        file_time = datetime.fromtimestamp(os.path.getmtime(cache_file))
        return (datetime.now() - file_time) < timedelta(hours=CACHE_EXPIRY_HOURS)

    def get_data(self, symbol, period_name):
        """
        获取股票K线数据
        :param symbol: 股票代码 (格式: 美股AAPL, A股600000.SH, 港股0700.HK)
        :param period_name: K线周期名称 (必须为SUPPORTED_PERIODS中的值)
        :return: 包含K线数据的DataFrame，失败时返回空DataFrame
        """
        # 验证周期名称
        if period_name not in SUPPORTED_PERIODS:
            logger.error(f"不支持的K线周期: {period_name}. 支持的周期: {SUPPORTED_PERIODS}")
            return pd.DataFrame()
        
        # 尝试使用缓存
        cache_file = self._get_cache_filename(symbol, period_name)
        if ENABLE_CACHE and self._is_cache_valid(cache_file):
            try:
                df = pd.read_csv(cache_file)
                df['timestamp'] = pd.to_datetime(df['timestamp'])
                logger.info(f"从缓存加载数据: {symbol} {period_name}")
                return df
            except Exception as e:
                logger.warning(f"缓存加载失败: {e}")
        
        # 🎯 智能数据源选择
        df = pd.DataFrame()
        attempts = 0

        # 检查是否为期货代码，如果是则强制使用AKShare
        is_futures = self._detect_futures_symbol(symbol)
        if is_futures:
            logger.info(f"🔥 检测到期货代码 {symbol}，强制使用AKShare数据源")
            effective_source = 'akshare'
        else:
            effective_source = self.data_source

        while attempts < MAX_RETRIES and df.empty:
            try:
                if effective_source == 'yfinance':
                    df = self._fetch_yfinance(symbol, period_name)
                elif effective_source == 'akshare':
                    df = self._fetch_akshare(symbol, period_name)
                elif effective_source == 'tushare' and TUSHARE_TOKEN:
                    df = self._fetch_tushare(symbol, period_name)
                elif effective_source == 'baostock':
                    df = self._fetch_baostock(symbol, period_name)
                else:
                    logger.error(f"未配置或未支持的数据源: {effective_source}")
                    return pd.DataFrame()
            except Exception as e:
                logger.error(f"获取数据失败(尝试 {attempts+1}/{MAX_RETRIES}): {e}")
                attempts += 1
                time.sleep(RETRY_DELAY)
        
        # 如果获取成功，保存到缓存
        if not df.empty and ENABLE_CACHE:
            try:
                df.to_csv(cache_file, index=False, encoding=OUTPUT_ENCODING)
                logger.info(f"数据已缓存至: {cache_file}")
            except Exception as e:
                logger.warning(f"缓存保存失败: {e}")
        
        return df

    def _detect_futures_symbol(self, symbol):
        """
        🎯 检测是否为期货代码
        用于智能选择数据源
        """
        if not symbol:
            return False

        import re
        symbol_upper = symbol.upper()

        # 期货代码特征模式
        futures_patterns = [
            r'^[A-Z]{1,3}\d{4}$',    # WR2510, RB2510等
            r'^[A-Z]{1,3}9999$',     # RB9999, WR9999等主力合约
            r'^[A-Z]{1,4}$'          # RB, WR等品种代码
        ]

        for pattern in futures_patterns:
            if re.match(pattern, symbol_upper):
                logger.info(f"🎯 识别期货代码模式: {symbol} -> {pattern}")
                return True

        return False

    def _fetch_yfinance(self, symbol, period_name):
        """
        从Yahoo Finance获取数据
        :param symbol: 股票代码
        :param period_name: K线周期名称
        :return: DataFrame(K线数据)
        """
        import yfinance as yf
        
        logger.info(f"从Yahoo Finance获取数据: {symbol} {period_name}")
        
        # 获取周期映射
        period = PERIOD_MAPPING.get('yfinance', {}).get(period_name, PERIOD_MAPPING.get(period_name, '1d'))
        
        # 根据周期确定时间范围
        if period_name == '60分钟':
            period = '60m'
            end_date = datetime.now()
            start_date = end_date - timedelta(days=60)  # 获取最近60天的数据
        elif period_name == '日线':
            start_date = datetime.now() - timedelta(days=365)  # 1年数据
            end_date = datetime.now()
        elif period_name == '周线':
            start_date = datetime.now() - timedelta(days=365 * 5)  # 5年数据
            end_date = datetime.now()
        elif period_name == '月线':
            start_date = datetime.now() - timedelta(days=365 * 10)  # 10年数据
            end_date = datetime.now()
        
        # 下载数据
        stock = yf.Ticker(symbol)
        df = stock.history(
            period='max' if period_name in ['周线', '月线'] else None,
            interval=period,
            start=start_date if period_name not in ['周线', '月线'] else None,
            end=end_date if period_name not in ['周线', '月线'] else None,
            timeout=API_TIMEOUT
        )
        
        # 处理数据
        if df.empty:
            logger.warning("未获取到数据")
            return pd.DataFrame()
        
        # 重置索引并重命名列
        df.reset_index(inplace=True)
        df.rename(columns={
            'Date': 'timestamp',
            'Open': 'open',
            'High': 'high',
            'Low': 'low',
            'Close': 'close',
            'Volume': 'volume'
        }, inplace=True)
        
        # 添加股票代码列
        df['symbol'] = symbol
        
        # 转换时间戳
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        
        logger.info(f"获取到 {len(df)} 条K线数据")
        return df

    def _fetch_akshare(self, symbol, period_name):
        """
        🚀 优化的AKShare数据获取方法
        根据官方文档 https://akshare.akfamily.xyz/ 进行优化配置
        :param symbol: 股票代码 (支持多种格式: 000001, 000001.SZ, 600000.SH等)
        :param period_name: K线周期名称
        :return: DataFrame(K线数据)
        """
        import akshare as ak

        logger.info(f"🔄 从AKShare获取数据: {symbol} {period_name}")

        try:
            # 🔧 智能股票代码处理
            clean_symbol = self._normalize_akshare_symbol(symbol)
            if not clean_symbol:
                logger.error(f"❌ 无法识别的股票代码格式: {symbol}")
                return pd.DataFrame()

            # 📅 智能时间范围计算
            start_date, end_date = self._calculate_date_range(period_name)

            # 📊 获取周期映射
            period_map = {
                '60分钟': 'daily',  # AKShare中分钟数据使用特殊接口
                '日线': 'daily',
                '周线': 'weekly',
                '月线': 'monthly'
            }
            period = period_map.get(period_name, 'daily')

            logger.info(f"📈 参数: 代码={clean_symbol}, 周期={period}, 时间={start_date}-{end_date}")

            # 🎯 智能选择数据接口（股票 vs 期货）
            df = pd.DataFrame()

            # 判断是股票还是期货
            is_futures = self._is_futures_symbol(clean_symbol)

            if is_futures:
                # 🔥 期货数据获取
                df = self._fetch_futures_data_akshare(clean_symbol, period_name, start_date, end_date)
            else:
                # 📈 股票数据获取
                df = self._fetch_stock_data_akshare(clean_symbol, period_name, start_date, end_date)

            # 📋 数据验证和清理
            if df.empty:
                logger.warning(f"⚠️ 未获取到 {symbol} 的数据")
                return pd.DataFrame()

            # 🔄 标准化数据格式
            df = self._standardize_akshare_data(df, symbol)

            logger.info(f"✅ 成功获取AKShare数据: {len(df)}条记录，时间范围: {df['timestamp'].min()} 至 {df['timestamp'].max()}")
            return df

        except Exception as e:
            logger.error(f"❌ AKShare数据获取失败 {symbol}: {e}")
            return pd.DataFrame()

    def _normalize_akshare_symbol(self, symbol):
        """
        🔧 标准化AKShare代码格式
        支持股票和期货代码: 000001, 000001.SZ, 600000.SH, WR2510, RB2510等
        """
        if not symbol:
            return None

        original_symbol = symbol.upper()

        # 1. 处理股票代码
        clean_symbol = original_symbol.replace('.SZ', '').replace('.SH', '').replace('.BJ', '')
        clean_symbol = clean_symbol.replace('SH', '').replace('SZ', '').replace('BJ', '')

        # 股票代码：6位数字
        if clean_symbol.isdigit() and len(clean_symbol) == 6:
            return clean_symbol

        # 2. 处理期货代码
        futures_info = self._identify_futures_code(original_symbol)
        if futures_info:
            return futures_info

        # 3. 处理主力合约（以9999结尾）
        if original_symbol.endswith('9999'):
            return original_symbol

        # 4. 处理简短代码（可能是期货品种）
        if len(original_symbol) <= 4 and original_symbol.isalpha():
            return original_symbol

        return None

    def _identify_futures_code(self, symbol):
        """
        🎯 识别期货代码格式
        支持格式: WR2510, RB2510, CU2510, AG2510等
        """
        if not symbol or len(symbol) < 4:
            return None

        # 期货代码模式：字母+数字（如WR2510, RB2510）
        import re

        # 模式1: 2-3个字母 + 4位数字 (如WR2510, RBU2510)
        pattern1 = r'^([A-Z]{1,3})(\d{4})$'
        match1 = re.match(pattern1, symbol)
        if match1:
            commodity = match1.group(1)
            contract = match1.group(2)
            logger.info(f"🎯 识别期货代码: {commodity} {contract}合约")
            return symbol

        # 模式2: 字母 + 数字 + 字母 + 数字 (如RBU2510)
        pattern2 = r'^([A-Z]{1,2})([A-Z]?)(\d{4})$'
        match2 = re.match(pattern2, symbol)
        if match2:
            logger.info(f"🎯 识别期货代码: {symbol}")
            return symbol

        # 模式3: 主力合约格式 (如RB9999)
        pattern3 = r'^([A-Z]{1,3})(9999)$'
        match3 = re.match(pattern3, symbol)
        if match3:
            logger.info(f"🎯 识别期货主力合约: {symbol}")
            return symbol

        return None

    def _is_futures_symbol(self, symbol):
        """
        🎯 判断是否为期货代码
        """
        if not symbol:
            return False

        # 期货代码特征
        import re

        # 1. 主力合约（以9999结尾）
        if symbol.endswith('9999'):
            return True

        # 2. 期货合约格式（字母+数字）
        futures_pattern = r'^[A-Z]{1,3}\d{4}$'
        if re.match(futures_pattern, symbol):
            return True

        # 3. 纯字母（期货品种代码）
        if symbol.isalpha() and len(symbol) <= 4:
            return True

        return False

    def _fetch_stock_data_akshare(self, clean_symbol, period_name, start_date, end_date):
        """
        📈 获取股票数据
        """
        import akshare as ak

        try:
            if period_name == '60分钟':
                # 使用分钟级数据接口
                df = ak.stock_zh_a_hist_min_em(
                    symbol=clean_symbol,
                    period='60',
                    adjust="qfq"  # 前复权
                )
                # 限制数据量（最近60天）
                if not df.empty and '时间' in df.columns:
                    df['时间'] = pd.to_datetime(df['时间'])
                    cutoff_date = datetime.now() - timedelta(days=60)
                    df = df[df['时间'] >= cutoff_date]
            else:
                # 使用日/周/月线数据接口
                period_map = {
                    '日线': 'daily',
                    '周线': 'weekly',
                    '月线': 'monthly'
                }
                period = period_map.get(period_name, 'daily')

                df = ak.stock_zh_a_hist(
                    symbol=clean_symbol,
                    period=period,
                    start_date=start_date,
                    end_date=end_date,
                    adjust="qfq"  # 前复权，避免除权除息影响
                )

            return df

        except Exception as e:
            logger.error(f"❌ 股票数据获取失败 {clean_symbol}: {e}")
            return pd.DataFrame()

    def _fetch_futures_data_akshare(self, clean_symbol, period_name, start_date, end_date):
        """
        🔥 获取期货数据
        """
        import akshare as ak

        try:
            logger.info(f"🔥 获取期货数据: {clean_symbol} {period_name}")

            # 期货品种代码映射
            futures_mapping = {
                'WR': 'WR',  # 线材
                'RB': 'RB',  # 螺纹钢
                'HC': 'HC',  # 热轧卷板
                'CU': 'CU',  # 铜
                'AL': 'AL',  # 铝
                'ZN': 'ZN',  # 锌
                'PB': 'PB',  # 铅
                'NI': 'NI',  # 镍
                'SN': 'SN',  # 锡
                'AU': 'AU',  # 黄金
                'AG': 'AG',  # 白银
                'RU': 'RU',  # 橡胶
                'BU': 'BU',  # 沥青
                'FU': 'FU',  # 燃料油
                'SC': 'SC',  # 原油
            }

            # 提取期货品种
            import re
            match = re.match(r'^([A-Z]{1,3})', clean_symbol)
            if match:
                commodity = match.group(1)
                logger.info(f"🎯 期货品种: {commodity}")

                # 🎯 优先使用历史数据接口（数据更完整）
                df = pd.DataFrame()

                # 方法1: 期货历史数据接口（优先）
                try:
                    df = ak.futures_zh_daily_sina(symbol=clean_symbol)
                    if not df.empty:
                        logger.info(f"✅ 使用历史数据接口获取 {clean_symbol} 数据: {len(df)} 条")
                    else:
                        logger.warning(f"⚠️ 历史数据接口返回空数据")
                except Exception as e1:
                    logger.warning(f"⚠️ 历史数据接口失败: {e1}")

                # 方法2: 如果历史数据失败，尝试现货接口
                if df.empty:
                    try:
                        df = ak.futures_zh_spot(symbol=clean_symbol)
                        if not df.empty:
                            logger.info(f"✅ 使用现货接口获取 {clean_symbol} 数据: {len(df)} 条")
                    except Exception as e2:
                        logger.warning(f"⚠️ 现货接口失败: {e2}")

                # 方法3: 主力合约接口
                if df.empty and clean_symbol.endswith('9999'):
                    try:
                        base_symbol = clean_symbol[:-4]
                        df = ak.futures_main_sina(symbol=base_symbol)
                        if not df.empty:
                            logger.info(f"✅ 使用主力合约接口获取 {base_symbol} 数据: {len(df)} 条")
                    except Exception as e3:
                        logger.warning(f"⚠️ 主力合约接口失败: {e3}")

                return df

            logger.warning(f"⚠️ 无法解析期货代码: {clean_symbol}")
            return pd.DataFrame()

        except Exception as e:
            logger.error(f"❌ 期货数据获取失败 {clean_symbol}: {e}")
            return pd.DataFrame()

    def _calculate_date_range(self, period_name):
        """
        📅 根据周期智能计算时间范围
        """
        end_date = datetime.now().strftime('%Y%m%d')

        # 根据周期优化数据获取范围
        date_ranges = {
            '60分钟': 60,      # 60天分钟数据
            '日线': 365 * 2,   # 2年日线数据
            '周线': 365 * 5,   # 5年周线数据
            '月线': 365 * 10   # 10年月线数据
        }

        days = date_ranges.get(period_name, 365)
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')

        return start_date, end_date

    def _standardize_akshare_data(self, df, original_symbol):
        """
        🔄 标准化AKShare数据格式
        """
        try:
            # 📋 完整的列名映射（支持股票和期货数据）
            column_mapping = {
                # 时间列
                '日期': 'timestamp', '时间': 'timestamp', 'datetime': 'timestamp',
                'date': 'timestamp', 'time': 'timestamp',

                # OHLC数据
                '开盘': 'open', 'open': 'open', '开盘价': 'open',
                '最高': 'high', 'high': 'high', '最高价': 'high',
                '最低': 'low', 'low': 'low', '最低价': 'low',
                '收盘': 'close', 'close': 'close', '收盘价': 'close',

                # 期货特有列名
                'current_price': 'close',  # 期货现价作为收盘价
                'last_close': 'prev_close',  # 昨收盘
                'settle': 'settle_price',    # 结算价
                'hold': 'open_interest',     # 持仓量

                # 成交量和成交额
                '成交量': 'volume', 'volume': 'volume', '成交量(股)': 'volume',
                '成交额': 'amount', 'amount': 'amount', '成交额(元)': 'amount',

                # 期货成交量相关
                'buy_vol': 'buy_volume',     # 买量
                'sell_vol': 'sell_volume',   # 卖量

                # 其他指标
                '振幅': 'amplitude', '涨跌幅': 'pct_change', '涨跌额': 'change',
                '换手率': 'turnover', '市盈率': 'pe_ratio', '市净率': 'pb_ratio',

                # 期货价格相关
                'bid_price': 'bid',          # 买价
                'ask_price': 'ask',          # 卖价
                'avg_price': 'avg_price',    # 均价
                'last_settle_price': 'prev_settle'  # 昨结算
            }

            # 应用列名映射
            df = df.rename(columns=column_mapping)

            # 🕐 处理时间列
            if 'timestamp' not in df.columns:
                # 尝试从索引获取时间
                if df.index.name in ['日期', 'date', 'time']:
                    df['timestamp'] = df.index
                else:
                    logger.error("❌ 无法找到时间列")
                    return pd.DataFrame()

            # 🕐 智能时间格式转换
            df['timestamp'] = self._convert_timestamp_smart(df['timestamp'])

            # 🔢 确保数值列为正确类型
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'amount']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 📊 数据质量检查
            # 移除价格为0或空值的行
            essential_columns = ['open', 'high', 'low', 'close']
            for col in essential_columns:
                if col in df.columns:
                    df = df[df[col] > 0]

            # 移除成交量为负的异常数据
            if 'volume' in df.columns:
                df = df[df['volume'] >= 0]

            # 🔄 按时间排序并重置索引
            df = df.sort_values('timestamp').reset_index(drop=True)

            # 📌 添加股票代码
            df['symbol'] = original_symbol

            # 🧹 移除重复数据
            df = df.drop_duplicates(subset=['timestamp'], keep='last')

            return df

        except Exception as e:
            logger.error(f"❌ 数据标准化失败: {e}")
            return df

    def _convert_timestamp_smart(self, timestamp_series):
        """
        🕐 智能时间格式转换
        处理各种时间格式，包括期货数据的特殊格式
        """
        try:
            # 如果已经是datetime类型，直接返回
            if pd.api.types.is_datetime64_any_dtype(timestamp_series):
                return timestamp_series

            # 处理数字格式的时间（如期货数据中的150000）
            if pd.api.types.is_numeric_dtype(timestamp_series):
                # 检查是否是时间格式（HHMMSS）
                sample_value = timestamp_series.iloc[0] if len(timestamp_series) > 0 else 0

                if 0 <= sample_value <= 240000:  # 可能是时间格式HHMMSS
                    logger.info(f"🕐 检测到时间格式数据，转换为今日时间: {sample_value}")
                    # 转换为今日的时间
                    today = datetime.now().date()
                    converted_times = []

                    for time_val in timestamp_series:
                        try:
                            # 将数字转换为时间字符串（补零到6位）
                            time_str = f"{int(time_val):06d}"
                            hour = int(time_str[:2])
                            minute = int(time_str[2:4])
                            second = int(time_str[4:6])

                            # 创建今日的datetime
                            dt = datetime.combine(today, datetime.min.time().replace(
                                hour=hour, minute=minute, second=second
                            ))
                            converted_times.append(dt)
                        except:
                            # 如果转换失败，使用当前时间
                            converted_times.append(datetime.now())

                    return pd.Series(converted_times, index=timestamp_series.index)

            # 尝试标准的pandas时间转换
            return pd.to_datetime(timestamp_series, errors='coerce')

        except Exception as e:
            logger.warning(f"⚠️ 时间转换失败: {e}，使用当前时间")
            # 如果所有转换都失败，返回当前时间的序列
            return pd.Series([datetime.now()] * len(timestamp_series), index=timestamp_series.index)

    def _fetch_tushare(self, symbol, period_name):
        """
        从Tushare获取数据（专业金融数据）
        :param symbol: 股票代码
        :param period_name: K线周期名称
        :return: DataFrame(K线数据)
        """
        if not TUSHARE_TOKEN:
            logger.error("Tushare token未配置")
            return pd.DataFrame()
        
        logger.info(f"从Tushare获取数据: {symbol} {period_name}")
        
        # 获取周期映射
        freq_map = {
            '60分钟': '60min',
            '日线': 'D',
            '周线': 'W',
            '月线': 'M'
        }
        freq = freq_map.get(period_name, 'D')

        # 确定时间范围
        end_date = datetime.now().strftime('%Y%m%d')
        if period_name == '60分钟':
            start_date = (datetime.now() - timedelta(days=60)).strftime('%Y%m%d')
        elif period_name == '日线':
            start_date = (datetime.now() - timedelta(days=365)).strftime('%Y%m%d')
        elif period_name == '周线':
            start_date = (datetime.now() - timedelta(days=365 * 5)).strftime('%Y%m%d')
        elif period_name == '月线':
            start_date = (datetime.now() - timedelta(days=365 * 10)).strftime('%Y%m%d')
        
        # 获取数据
        if symbol.startswith(('00', '30', '60')):
            # A股代码
            ts_code = symbol + '.SH' if symbol.startswith(('60', '68')) else symbol + '.SZ'
            if period_name == '60分钟':
                df = self.tushare_pro.btcbar(
                    ts_code=ts_code,
                    freq='60min',
                    start_date=start_date,
                    end_date=end_date
                )
            else:
                df = self.tushare_pro.daily(
                    ts_code=ts_code,
                    start_date=start_date,
                    end_date=end_date
                )
        else:
            logger.error(f"Tushare暂不支持非A股数据: {symbol}")
            return pd.DataFrame()
        
        # 处理数据
        if df.empty:
            logger.warning("未获取到数据")
            return pd.DataFrame()
        
        # 重命名列
        df.rename(columns={
            'trade_date': 'timestamp',
            'open': 'open',
            'high': 'high',
            'low': 'low',
            'close': 'close',
            'vol': 'volume',
            'amount': 'amount'
        }, inplace=True)
        
        # 添加股票代码列
        df['symbol'] = symbol
        
        # 转换时间戳
        df['timestamp'] = pd.to_datetime(df['timestamp'], format='%Y%m%d')
        
        # 排序数据
        df.sort_values('timestamp', inplace=True)
        
        logger.info(f"获取到 {len(df)} 条K线数据")
        return df

    def _fetch_baostock(self, symbol, period_name):
        """
        从Baostock获取数据（免费A股数据平台）
        :param symbol: 股票代码
        :param period_name: K线周期名称
        :return: DataFrame(K线数据)
        """
        logger.info(f"从Baostock获取数据: {symbol} {period_name}")

        # 登录Baostock系统
        lg = self.bs.login()
        if lg.error_code != '0':
            logger.error(f"Baostock登录失败: {lg.error_msg}")
            return pd.DataFrame()

        try:
            # 转换股票代码格式
            bs_code = self._convert_to_baostock_code(symbol)
            if not bs_code:
                logger.error(f"不支持的股票代码格式: {symbol}")
                return pd.DataFrame()

            # 获取周期映射
            freq_map = {
                '60分钟': '60',  # 60分钟
                '日线': 'd',    # 日线
                '周线': 'w',    # 周线
                '月线': 'm'     # 月线
            }
            frequency = freq_map.get(period_name, 'd')

            # 确定时间范围
            end_date = datetime.now().strftime('%Y-%m-%d')
            if period_name == '60分钟':
                start_date = (datetime.now() - timedelta(days=60)).strftime('%Y-%m-%d')
            elif period_name == '日线':
                start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')
            elif period_name == '周线':
                start_date = (datetime.now() - timedelta(days=365 * 5)).strftime('%Y-%m-%d')
            elif period_name == '月线':
                start_date = (datetime.now() - timedelta(days=365 * 10)).strftime('%Y-%m-%d')

            # 获取K线数据
            if period_name == '60分钟':
                # 获取分钟级数据
                rs = self.bs.query_history_k_data_plus(
                    bs_code,
                    "date,time,code,open,high,low,close,volume,amount",
                    start_date=start_date,
                    end_date=end_date,
                    frequency=frequency,
                    adjustflag="3"  # 不复权
                )
            else:
                # 获取日线及以上级别数据
                # 周线和月线不支持某些字段，使用基础字段
                if period_name in ['周线', '月线']:
                    fields = "date,code,open,high,low,close,volume,amount"
                else:
                    fields = "date,code,open,high,low,close,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"

                rs = self.bs.query_history_k_data_plus(
                    bs_code,
                    fields,
                    start_date=start_date,
                    end_date=end_date,
                    frequency=frequency,
                    adjustflag="3"  # 不复权
                )

            if rs.error_code != '0':
                logger.error(f"Baostock数据获取失败: {rs.error_msg}")
                return pd.DataFrame()

            # 转换为DataFrame
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                logger.warning("未获取到数据")
                return pd.DataFrame()

            # 创建DataFrame
            if period_name == '60分钟':
                columns = ["date", "time", "code", "open", "high", "low", "close", "volume", "amount"]
            elif period_name in ['周线', '月线']:
                columns = ["date", "code", "open", "high", "low", "close", "volume", "amount"]
            else:
                columns = ["date", "code", "open", "high", "low", "close", "volume", "amount",
                          "adjustflag", "turn", "tradestatus", "pctChg", "isST"]

            df = pd.DataFrame(data_list, columns=columns)

            # 处理时间戳
            if period_name == '60分钟':
                # 处理Baostock的时间格式：date='2025-05-20', time='20250520103000000'
                # 提取时间部分：从time字段中提取HHMMSS
                def parse_baostock_time(date_str, time_str):
                    try:
                        # time格式: 20250520103000000，提取中间的HHMMSS部分
                        if len(time_str) >= 14:
                            time_part = time_str[8:14]  # 提取HHMMSS
                            hour = time_part[:2]
                            minute = time_part[2:4]
                            second = time_part[4:6]
                            return f"{date_str} {hour}:{minute}:{second}"
                        else:
                            return f"{date_str} 00:00:00"
                    except:
                        return f"{date_str} 00:00:00"

                df['timestamp'] = df.apply(lambda row: parse_baostock_time(row['date'], row['time']), axis=1)
                df['timestamp'] = pd.to_datetime(df['timestamp'])
            else:
                df['timestamp'] = pd.to_datetime(df['date'])

            # 转换数据类型
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'amount']
            for col in numeric_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

            # 添加股票代码列
            df['symbol'] = symbol

            # 选择需要的列
            df = df[['timestamp', 'open', 'high', 'low', 'close', 'volume', 'symbol']]

            # 过滤无效数据
            df = df.dropna(subset=['open', 'high', 'low', 'close'])
            df = df[df['volume'] > 0]  # 过滤成交量为0的数据

            # 排序数据
            df.sort_values('timestamp', inplace=True)
            df.reset_index(drop=True, inplace=True)

            logger.info(f"获取到 {len(df)} 条K线数据")
            return df

        except Exception as e:
            logger.error(f"Baostock数据处理失败: {e}")
            return pd.DataFrame()
        finally:
            # 登出Baostock系统
            self.bs.logout()

    def _convert_to_baostock_code(self, symbol):
        """
        转换股票代码为Baostock格式
        :param symbol: 原始股票代码
        :return: Baostock格式的股票代码
        """
        # 处理不同格式的股票代码
        if '.' in symbol:
            # 格式: 600000.SH, 000001.SZ
            code, exchange = symbol.split('.')
            if exchange.upper() == 'SH':
                return f"sh.{code}"
            elif exchange.upper() == 'SZ':
                return f"sz.{code}"
        else:
            # 纯数字代码，根据开头判断交易所
            if symbol.startswith(('60', '68', '90')):
                return f"sh.{symbol}"
            elif symbol.startswith(('00', '30', '20')):
                return f"sz.{symbol}"

        logger.warning(f"无法识别的股票代码格式: {symbol}")
        return None

    def save_to_file(self, df, symbol, period_name):
        """
        保存数据到文件
        :param df: 包含数据的DataFrame
        :param symbol: 股票代码
        :param period_name: K线周期名称
        :return: 保存的文件路径
        """
        if df.empty:
            logger.warning("数据为空，不保存")
            return None
        
        # 生成文件名
        date_str = datetime.now().strftime("%Y%m%d")
        filename = OUTPUT_FILENAME_FORMAT.format(
            symbol=symbol,
            period=period_name,
            date=date_str
        )
        
        try:
            # 确保数据按时间排序
            df.sort_values('timestamp', inplace=True)
            
            # 保存到CSV
            df.to_csv(filename, index=False, encoding=OUTPUT_ENCODING)
            logger.info(f"数据已保存至: {filename}")
            return filename
        except Exception as e:
            logger.error(f"保存数据失败: {e}")
            return None

# 测试函数
def test_data_fetcher():
    """测试数据获取功能"""
    print("=" * 50)
    print("数据获取模块测试")
    print("=" * 50)
    
    fetcher = DataFetcher()
    
    # 测试股票
    test_cases = [
        # (股票代码, 周期)
        ('AAPL', '60分钟'),  # 苹果股票60分钟线
        ('600000.SH', '日线'),  # 浦发银行日线
        ('0700.HK', '周线'),  # 腾讯周线
        ('RB9999', '月线'),  # 螺纹钢主力合约月线
    ]
    
    for symbol, period in test_cases:
        print(f"\n获取 {symbol} 的 {period} K线数据...")
        start_time = time.time()
        df = fetcher.get_data(symbol, period)
        elapsed = time.time() - start_time
        
        if df.empty:
            print(f"获取失败: {symbol} {period}")
            continue
            
        print(f"成功获取 {len(df)} 条数据, 耗时: {elapsed:.2f}秒")
        print(f"时间范围: {df['timestamp'].min()} 到 {df['timestamp'].max()}")
        print(f"列名: {list(df.columns)}")
        print(f"前2条数据:\n{df.head(2)}")
        
        # 保存文件
        filename = fetcher.save_to_file(df, symbol, period)
        if filename:
            print(f"已保存到: {filename}")

if __name__ == "__main__":
    # 运行测试
    test_data_fetcher()